# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com

from FiniteDifferencesEngineMod import *
from copy import deepcopy
from numpy import interp
from Tools import multiDimensionalLinearInterpolation
from math import pow
from MatrixSolvers import *

class ImplicitFiniteDifferencesEngine(FiniteDifferencesEngine):

    def initialise(self, relations, initials):
        self.transientIndicesToCalculate = []
        self.relations = relations
        for relation in self.relations:
            for varName in relation.variableNames:
                relation.varIndices.append(self.grid.findIndexInNamesOrder(varName))

        #   first, find the 'end' node
        self.firstCoord = [0 for idx in range(len(initials))]
        self.distances = [0 for idx in range(len(initials))]
        for idx in range(len(initials)):
            gridIdx = self.grid.findIndexInNamesOrder(initials.keys()[idx])
            for jdx in range(self.grid.dimSizes[gridIdx]):
                if self.grid.g[gridIdx][jdx] >= initials.values()[idx]:
                    self.firstCoord[gridIdx] = jdx
                    self.distances[gridIdx] = self.grid.g[gridIdx][jdx] - initials.values()[idx]
                    break

        isForwards = self.firstCoord[self.transientIndex] == self.grid.dimSizes[self.transientIndex] - 1
        if not isForwards:self.transientIndicesToCalculate
            if self.firstCoord[self.transientIndex] == 0:
                self.transientIndicesToCalculate = range(self.grid.dimSizes[self.transientIndex])
            else:
                raise Exception('Cannot determine transient type')
        else:
            self.transientIndicesToCalculate = range(self.grid.dimSizes[self.transientIndex]-1,-1,-1)

        #   determine the algortihm for the matrix to be solved, i.e. tridiagonal for 1d pde, etc
        self.solver = ''
        if len(self.relations) == 3:
            self.solver = 'solveTridiagonal'
        elif len(self.relations) > 3:
            self.solver = 'solveGaussSiedel'
        self.matrix = []
        variableIndices = [i for i in range(self.grid.ndims)]
        variableIndices.remove(self.transientIndex)
        self.matrixLength = 1.0
        for idx in range(len(variableIndices)):
            self.matrixLength *= self.grid.dimSizes[variableIndices[idx]]
        for relation in self.relations:
            self.matrix.append([0.0 for i in range(self.matrixLength)])
            

    def doCalculation(self):
        for tIndex in self.transientIndicesToCalculate :
            # build the matrix to be inverted
            self.matrix
            for relation in self.relations: